import glob
import torch
import warnings
from multiprocessing import cpu_count
from .base_options import BaseOptions
from options import *
from utils.utils import *

class TrainOptions(BaseOptions):
    """模型训练参数选项"""

    def initialize(self):
        parser = BaseOptions.initialize(self)
        parser.add_argument("--dataset", required=True, help="训练数据集路径（此路径下应至少包含一个子文件夹）")
        parser.add_argument("--name", type=str, default=None, help="存放训练数据的文件夹名称，默认为风格图片的名称")
        parser.add_argument("--style_image", type=str, required=True, help="风格图片的路径")
        parser.add_argument("--content_size", type=int, default=256, help="训练图片的大小，默认为 256x256")
        parser.add_argument("--style_size", type=int, default=None, help="风格图片的大小，默认为原图大小")
        parser.add_argument("--batch_size", type=int, default=16, help="每一次训练迭代中数据的批量大小")
        parser.add_argument("--num_workers", type=int, default=4, help="并行读取数据的进程个数")
        parser.add_argument("--arch", type=str, default="vgg16", help="作图片特征提取的网络名称，目前支持的有 vgg16, vgg19")
        parser.add_argument("--content_layers", type=str, default="8", help="图片内容特征提取的层的索引序列")
        parser.add_argument("--style_layers", type=str, default="3,8,15,22", help="图片风格特征提取的层的索引序列")
        parser.add_argument("--content_weight", type=float, default=1e5, help="内容损失的权重")
        parser.add_argument("--style_weight", type=float, default=1e10, help="风格损失的权重")
        parser.add_argument("--lr", type=float, default=1e-3, help="学习率")
        parser.add_argument("--gpu_ids", type=str, default="0", help="GPU 设备的 id 序列，如 0,1，不使用 GPU 应设置为 -1")
        parser.add_argument("--num_epochs", type=int, default=1, help="模型训练所迭代的次数")
        parser.add_argument("--epoch_count", type=int, default=0, help="模型起始迭代的 epoch，默认从 0 开始")
        parser.add_argument("--checkpoints_dir", type=str, default="./checkpoints", help="checkpoints 跟路径")
        parser.add_argument("--checkpoint_model", type=str, default=None, help="恢复训练 checkpoints 的模型路径")
        parser.add_argument("--checkpoint_interval", type=int, default=100, help="训练模型保存周期（每迭代次数）")
        parser.add_argument("--log_interval", type=int, default=1, help="日志输出周期（每迭代次数）")
        parser.add_argument("--sample_interval", type=int, default=50, help="随机样本测试周期（每迭代次数）")
        return parser

    def parse(self):
        opt = self.initialize().parse_args()

        # 如未指定 style_image 由 name 寻找风格图片
        if is_empty(opt.style_image):
            assert not is_empty(opt.name), "'style_image' 和 'name' 两个参数不能同时为空"
            style_images = glob.glob(f"images/style/{opt.name}.*")
            assert not is_empty(style_images), f"默认路径下 '{opt.name}' 所指定的风格图片不存在"
            opt.style_image = style_images[0]  # 默认情况下取第一张作为风格图

        # 初始化工作空间文件夹名称
        if is_empty(opt.name):
            opt.name = parse_filename(opt.style_image, '/')

        # 检查并行读取数据的进程个数
        supported_cores = cpu_count()
        if int(opt.num_workers) > supported_cores:
            warnings.warn(f"并行读取数据的进程个数 {opt.num_workers} 大于 CPU 支持的核心个数 {supported_cores}")

        # 内容特征和风格特征层索引转换
        opt.style_layers = parse_sequence(opt.style_layers, type_func=int)
        opt.content_layers = parse_sequence(opt.content_layers, type_func=int)

        # 检查训练使用的 GPU 设备序列
        supported_gpu_ids = list(range(torch.cuda.device_count())) + [-1]
        opt.gpu_ids = [-1] if opt.gpu_ids == "-1" else parse_sequence(opt.gpu_ids, type_func=int)
        assert set(opt.gpu_ids).issubset(supported_gpu_ids), \
            f"非法设备 id，当前设备支持的 GPU 序列有 {', '.join(map(lambda x: str(x), supported_gpu_ids))}"

        self.print_options(opt)
        return opt
